// 堆 https://blog.csdn.net/rxbook/article/details/130996393
package main

import "fmt"

// 堆结构体
type Heap struct {
	data []int //存储堆元素的数组,下标从1开始
	cap  int   //堆中可以存储的最大元素个数(容量)
	len  int   //堆中已经存储的元素个数(长度)
}

// 初始化一个堆
func NewHeap(capacity int) *Heap {
	heap := &Heap{}
	heap.cap = capacity
	heap.data = make([]int, capacity+1)
	heap.len = 0
	return heap
}

// 大顶堆插入元素,从下往上堆化
func (heap *Heap) insert(data int) {
	//判断堆是否已经满了
	if heap.len == heap.cap {
		return
	}
	heap.len++
	heap.data[heap.len] = data

	//和父节点互相比较大小
	i := heap.len
	parent := i / 2
	for parent > 0 && heap.data[parent] < heap.data[i] { //元素上浮,从下往上堆化
		swap(heap.data, parent, i) //和父节点的元素交换位置
		i = parent
		parent = i / 2
	}
}

// 大顶堆删除元素,从上往下堆化
func (heap *Heap) removeMax() {
	//判断堆是否已经满了
	if heap.len == 0 {
		return
	}

	//交换元素
	swap(heap.data, 1, heap.len)
	heap.data[heap.len] = 0
	heap.len--

	//元素下沉,从上往下堆化
	for i := 1; i <= heap.len/2; {
		maxIndex := i
		if heap.data[i] < heap.data[i*2] {
			maxIndex = i * 2
		}
		if i*2+1 <= heap.len && heap.data[maxIndex] < heap.data[i*2+1] {
			maxIndex = i*2 + 1
		}
		if maxIndex == i {
			break
		}
		swap(heap.data, i, maxIndex)
		i = maxIndex
	}
}

// 交换数组中两个元素
func swap(a []int, i int, j int) {
	a[i], a[j] = a[j], a[i]
}

func main() {
	//初始化堆,并插入6个元素
	heap := NewHeap(7)
	heap.insert(3)
	heap.insert(4)
	heap.insert(6)
	heap.insert(7)
	heap.insert(8)
	heap.insert(9)
	fmt.Println(heap) //&{[0 9 7 8 3 6 4 0] 7 6}

	//给堆中插入一个新元素 10
	heap.insert(10)
	fmt.Println(heap) //&{[0 10 7 9 3 6 4 8] 7 7}

	//移除堆中最大的元素
	heap.removeMax()
	fmt.Println(heap) //&{[0 9 7 8 3 6 4 0] 7 6}
}
